// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"log"
	"os"

	"google.golang.org/genai"

	"google.golang.org/adk/agent"
	"google.golang.org/adk/agent/llmagent"
	"google.golang.org/adk/agent/workflowagents/sequentialagent"
	"google.golang.org/adk/cmd/launcher"
	"google.golang.org/adk/cmd/launcher/full"
	"google.golang.org/adk/model/gemini"
)

func main() {
	ctx := context.Background()

	model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{})
	if err != nil {
		log.Fatalf("failed to create model: %s", err)
	}

	// --- 1. Define Sub-Agents for Each Pipeline Stage ---

	// Code Writer Agent
	// Takes the initial specification (from user query) and writes code.
	codeWriterAgent, err := llmagent.New(llmagent.Config{
		Name:  "CodeWriterAgent",
		Model: model,
		Instruction: `You are a Python Code Generator.
Based *only* on the user's request, write Python code that fulfills the requirement.
Output *only* the complete Python code block, enclosed in triple backticks ('''python ... ''').
Do not add any other text before or after the code block.`,
		Description: "Writes initial Python code based on a specification.",
		OutputKey:   "generated_code", // Stores output in state["generated_code"]
	})
	if err != nil {
		log.Fatalf("failed to create codeWriterAgent: %s", err)
	}

	// Code Reviewer Agent
	// Takes the code generated by the previous agent (read from state) and provides feedback.
	codeReviewerAgent, err := llmagent.New(llmagent.Config{
		Name:  "CodeReviewerAgent",
		Model: model,
		Instruction: `You are an expert Python Code Reviewer.
Your task is to provide constructive feedback on the provided code.

**Code to Review:**
'''python
{generated_code}
'''

**Review Criteria:**
1.  **Correctness:** Does the code work as intended? Are there logic errors?
2.  **Readability:** Is the code clear and easy to understand? Follows PEP 8 style guidelines?
3.  **Efficiency:** Is the code reasonably efficient? Any obvious performance bottlenecks?
4.  **Edge Cases:** Does the code handle potential edge cases or invalid inputs gracefully?
5.  **Best Practices:** Does the code follow common Python best practices?

**Output:**
Provide your feedback as a concise, bulleted list. Focus on the most important points for improvement.
If the code is excellent and requires no changes, simply state: "No major issues found."
Output *only* the review comments or the "No major issues" statement.`,
		Description: "Reviews code and provides feedback.",
		OutputKey:   "review_comments", // Stores output in state["review_comments"]
	})
	if err != nil {
		log.Fatalf("failed to create codeReviewerAgent: %s", err)
	}

	// Code Refactorer Agent
	// Takes the original code and the review comments (read from state) and refactors the code.
	codeRefactorerAgent, err := llmagent.New(llmagent.Config{
		Name:  "CodeRefactorerAgent",
		Model: model,
		Instruction: `You are a Python Code Refactoring AI.
Your goal is to improve the given Python code based on the provided review comments.

**Original Code:**
'''python
{generated_code}
'''

**Review Comments:**
{review_comments}

**Task:**
Carefully apply the suggestions from the review comments to refactor the original code.
If the review comments state "No major issues found," return the original code unchanged.
Ensure the final code is complete, functional, and includes necessary imports and docstrings.

**Output:**
Output *only* the final, refactored Python code block, enclosed in triple backticks ('''python ... ''').
Do not add any other text before or after the code block.`,
		Description: "Refactors code based on review comments.",
		OutputKey:   "refactored_code", // Stores output in state["refactored_code"]
	})
	if err != nil {
		log.Fatalf("failed to create codeRefactorerAgent: %s", err)
	}

	// --- 2. Create the SequentialAgent ---
	// This agent orchestrates the pipeline by running the sub_agents in order.
	codePipelineAgent, err := sequentialagent.New(sequentialagent.Config{
		AgentConfig: agent.Config{
			Name: "CodePipelineAgent",
			SubAgents: []agent.Agent{
				codeWriterAgent,
				codeReviewerAgent,
				codeRefactorerAgent,
			},
			Description: "Executes a sequence of code writing, reviewing, and refactoring.",
		},
	})
	if err != nil {
		log.Fatalf("failed to create codePipelineAgent: %s", err)
	}

	rootAgent := codePipelineAgent

	// The rootAgent can now be used by the ADK framework.
	log.Printf("Successfully created root agent: %s", rootAgent.Name())

	config := &launcher.Config{
		AgentLoader: agent.NewSingleLoader(rootAgent),
	}
	l := full.NewLauncher()
	if err = l.Execute(ctx, config, os.Args[1:]); err != nil {
		log.Fatalf("Run failed: %v\n\n%s", err, l.CommandLineSyntax())
	}
}
